##这个是10s拍一次照片识别
from ultralytics import YOLO
import cv2
import time
import os
from PIL import Image

# 创建输出文件夹
output_folder = 'imgsOut'
input_folder = 'imgsIn'
os.makedirs(output_folder, exist_ok=True)
os.makedirs(input_folder, exist_ok=True)

# 加载预训练的YOLOv8n模型
model = YOLO('results_runs/dataSet_good2/E10_IZ64/classify/train/weights/best.pt')

# 从摄像头捕捉图像
cap = cv2.VideoCapture(0)  # 0表示默认摄像头

while True:
    ret, frame = cap.read()  # 读取图像帧
    if not ret:
        break

    # 在图像上运行推理
    results = model(frame)

    # 获取当前时间戳
    timestamp = int(time.time())

    # 保存图像到输入文件夹
    input_image_path = os.path.join(input_folder, f'{timestamp}.jpg')
    cv2.imwrite(input_image_path, frame)

    # 处理识别结果
    for i, r in enumerate(results):
        im_array = r.plot()  # 绘制包含预测结果的BGR numpy数组
        im = Image.fromarray(im_array[..., ::-1])  # RGB PIL图像
        output_image_path = os.path.join(output_folder, f'{timestamp}_{i}.jpg')
        im.save(output_image_path)  # 保存识别结果图像

    time.sleep(10)  # 每隔10秒拍摄一次图像

cap.release()  # 释放摄像头
cv2.destroyAllWindows()  # 关闭所有窗口
